{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "79ebfce1-f5ac-4e39-83db-11416e310e8e",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Jupyter AI\n",
    "\n",
    "Jupyter AI is a JupyterLab extension that provides a friendly user interface between users and AI models.\n",
    "\n",
    "This demo showcases the IPython magics Jupyter AI provides out-of-the-box.\n",
    "\n",
    "Load the IPython extension:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "24f3f446-2b1d-4802-a47c-d298c06fc86e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%load_ext jupyter_ai"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a057ea12-a45c-4ab4-a300-82496d96df4e",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Getting started\n",
    "\n",
    "To start using Jupyter AI, use the `%%ai` cell magic with the model specified via the syntax `<provider-id>:<model-id>`. Every line after the first should contain the prompt.\n",
    "\n",
    "Currently, we support the following providers:\n",
    "\n",
    "- `ai21`\n",
    "- `anthropic`\n",
    "- `cohere`\n",
    "- `huggingface_hub`\n",
    "- `openai`\n",
    "- `openai-chat`\n",
    "- `sagemaker-endpoint`\n",
    "\n",
    "\n",
    "These are LangChain LLM providers (`langchain.llms`). Any model ID that is valid for the corresponding provider is valid in the AI magics. For example, `claude-v1.2` is a valid model ID for the `anthropic` provider, and one can invoke it via Jupyter AI magics like so:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0887d627-ff84-4042-86d6-90ac18a4f391",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "\n",
       "O C++, language so keen,\n",
       "\n",
       " With classes, objects and templates so sheen.\n",
       "\n",
       " Inheritance and polymorphism galore,\n",
       "\n",
       "Abstracting the complexity to the core.\n",
       "\n",
       " Support for low and high level alike,\n",
       "\n",
       "Helping programmers day and night.\n",
       "\n",
       "Pointers, references and RAII,\n",
       "\n",
       "OOP and generic programming, my oh my!\n",
       "\n",
       " Exception handling and lambda's too,\n",
       "\n",
       "There's so much you help me do! \n",
       "\n",
       " A powerful language, sometimes complex,\n",
       "\n",
       "But in the end, truly adept.\n",
       "\n",
       " That's my poem about C++,\n",
       "\n",
       " A language with power, through and through!"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%ai anthropic:claude-v1.2\n",
    "Write a poem about C++."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a565a33-05a9-4e4d-b204-768d572d561a",
   "metadata": {},
   "source": [
    "## Implicit providers\n",
    "\n",
    "If your model ID is unique, then we can also infer the provider ID implicitly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "976934ed-90aa-47db-9a69-843565eae02f",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "Copy code\n",
       "console\n",
       ".\n",
       "log\n",
       "(\n",
       "\"hello\n",
       "world\"\n",
       ");"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%ai j2-jumbo-instruct\n",
    "Write some JavaScript code that prints \"hello world\" to the console."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e178db2-55f7-45c4-a63c-9892ddf61495",
   "metadata": {},
   "source": [
    "## Rapidly experimenting with different HF Hub models\n",
    "We can call models on Hugging Face Hub directly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0837ad88-7759-42e8-894c-82747c4bb3ce",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "\n",
       "If that's the question you're looking for, here's the bottom line for New York: $2.13 per day.\n",
       "\n",
       "It's not surprising then, because as you can see"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%ai huggingface_hub:google/flan-t5-xl\n",
    "What is the capital of New York state?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "31f3e6e3-48cf-4e60-96d3-8b8e1dd34bec",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "\n",
       "2 is the number of times a point in a circle is square. Thus, one is squared twice: 2 x 2 = 19 x 18 = 51 / 2 x 2 = 49, and I was"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%ai huggingface_hub:gpt2\n",
    "What is the square root of 2?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a84c5776-0a57-4631-a441-056a707ceaf1",
   "metadata": {},
   "source": [
    "## Formatting via `-f/--format`\n",
    "\n",
    "You can also pass the `-f/--format` argument to specify an IPython display to use to render the output. Valid formats:\n",
    "\n",
    "- `markdown`\n",
    "- `math`\n",
    "- `html`\n",
    "- `json`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "190d7fd7-5359-430a-bccd-63d6f33559a2",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<svg width=\"200\" height=\"200\">\n",
       "  <rect x=\"50\" y=\"50\" width=\"100\" height=\"100\" style=\"fill: white; stroke: black; stroke-width: 5;\" /> \n",
       "</svg>\n",
       "\n",
       "This will create a 100x100 square with a 5px black border and white fill. \n",
       "Let me know if you have any other questions!"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%ai anthropic:claude-v1.2 -f html\n",
    "Create a square using SVG with a black border and white fill. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "0ddc7aff-d7d7-41dd-9371-4d636fa22d40",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\partial u}{\\partial t} =\\frac{\\partial^2 u}{\\partial x^2} + \\frac{\\partial^2 u}{\\partial y^2}$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%ai chatgpt -f math\n",
    "Generate the 2D heat equation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67fc5b32-92cf-4925-b1de-28617959e477",
   "metadata": {},
   "source": [
    "## IPython interpolation\n",
    "\n",
    "You can also interpolate IPython scope into your prompt via curly braces."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "be568759-5b76-4278-8201-e7fc61befc2f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "poet = \"Walt Whitman\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "36e60dca-68e3-49c7-bde4-25b7ea53eae4",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "O Captain, my Captain, you have led us through the fray  \n",
       "Through the storm and the turmoil, through the darkest of days  \n",
       "With your unwavering courage, your steadfast heart ablaze  \n",
       "You have lifted us up where the eagles soar and play  \n",
       "\n",
       "Your voice is as thunder, your soul a raging fire  \n",
       "You guide us through the tempest, you lift us ever higher  \n",
       "With you at our helm, we shall never falter or tire  \n",
       "For you are the Captain, and we are your loyal choir  \n",
       "\n",
       "Together we shall march, to the beat of our own drum  \n",
       "We shall conquer every fear, we shall overcome  \n",
       "With our eyes on the prize, and our hearts full of song  \n",
       "We shall claim victory, as we journey ever on  \n",
       "\n",
       "So here's to you, O Captain, our light in the dark  \n",
       "May your courage never falter, may your fire never spark  \n",
       "For in your hands we trust, our fate and our ark  \n",
       "O Captain, my Captain, you are the beating of our heart."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%ai chatgpt\n",
    "Write a poem in the style of {poet}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "54e72526-4f79-406a-a070-01c4373bc84a",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "1\n",
      "2\n",
      "3\n",
      "4\n"
     ]
    }
   ],
   "source": [
    "for i in range(0, 5):\n",
    "  print(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "9dc5f070-6656-4170-9c70-56484d1fe258",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "\n",
       "The above code will print the numbers 0 through 4. It will not print the number 5 because the range function includes the first value (i.e. 0) but not the last value (i.e. 5)."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%ai cohere:command-xlarge-nightly\n",
    "Please explain the code below:\n",
    "--\n",
    "{In[11]}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de1ed3fd-e9a0-43e5-8968-4f5001cb535d",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Full example: The Laplace Equation\n",
    "\n",
    "Say I just got a homework assignment about \"the Laplace equation\" from my professor.\n",
    "\n",
    "1. What is the Laplace equation?\n",
    "2. What is the form of the Laplace equation in polar coordinates?\n",
    "3. Write a Python algorithm to solve the Laplace equation numerically.\n",
    "\n",
    "The problem is: I have absolutely no idea what they're talking about.\n",
    "\n",
    "Let's ask an LLM what the assignment is about."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "28044123-5414-4123-ab26-6638ac47673a",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "\n",
       "The 2D Laplace equation is: \n",
       "$$ \\nabla^2 f = \\frac{\\partial^2 f}{\\partial x^2} + \\frac{\\partial^2 f}{\\partial y^2} = 0 $$ \n",
       "where $f$ is a scalar function. \n",
       "\n",
       "It is used to describe steady-state heat flow, electrostatics, fluid flow, and other phenomena in two dimensions under the assumption of no sources. The Laplace equation arises from simplifying the Poisson equation by setting the charge density to 0. It is a key partial differential equation used in modeling many physical systems.\n",
       "\n",
       "The Laplace equation implies that the sum of second derivatives of $f$ with respect to $x$ and $y$ is zero, which means the curvature of the solution surface is constant. So, the solutions have a minimal surface with constant negative Gaussian curvature.\n",
       "\n",
       "Does this help explain the 2D Laplace equation? Let me know if you have any other questions!"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%ai claude-v1.2\n",
    "What is the 2d Laplace equation and when is it used? Use LaTeX for equations delimited by `$`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "9a2c0af2-f1b5-4307-912a-3fc859ada08f",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \n",
       "\\frac{\\partial^2 \\psi}{\\partial r^2} + \\frac{1}{r} \\frac{\\partial \\psi}{\\partial r} + \\frac{1}{r^2} \\frac{\\partial^2 \\psi}{\\partial \\theta^2} = f\n",
       "$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%ai j2-jumbo-instruct --format math\n",
    "Write the 2d Laplace equation in polar coordinates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "3acfdd0d-830f-4281-959b-7e86e48610f1",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "Here's the Python code to solve the 2D Laplace equation in Cartesian coordinates:\n",
       "\n",
       "```python\n",
       "import numpy as np\n",
       "import matplotlib.pyplot as plt\n",
       "\n",
       "# Set up grid\n",
       "nx = 101\n",
       "ny = 101\n",
       "nt = 100\n",
       "\n",
       "dx = 1. / (nx - 1)\n",
       "dy = 1. / (ny - 1)\n",
       "\n",
       "x = np.linspace(0, 1, nx)\n",
       "y = np.linspace(0, 1, ny)\n",
       "\n",
       "# Initialize solution\n",
       "u = np.zeros((nx, ny))\n",
       "\n",
       "# Set boundary conditions\n",
       "u[0,:] = 0\n",
       "u[-1,:] = 0\n",
       "u[:,0] = 0\n",
       "u[:,-1] = 0\n",
       "\n",
       "# Set up source term\n",
       "f = np.zeros((nx, ny))\n",
       "f[int(nx/2), int(ny/2)] = 1 / (dx * dy)\n",
       "\n",
       "# Set up iteration parameters\n",
       "tolerance = 1e-5\n",
       "omega = 1.8\n",
       "\n",
       "# Iterative solver\n",
       "for n in range(nt):\n",
       "\n",
       "    for i in range(1, nx-1):\n",
       "        for j in range(1, ny-1):\n",
       "            u_old = u[i,j]\n",
       "            u[i,j] = (1-omega)*u[i,j] + omega*(dx**2*(u[i+1,j]+u[i-1,j]) + dy**2*(u[i,j+1]+u[i,j-1]) - dx**2*dy**2*f[i,j]) / (2*(dx**2+dy**2))\n",
       "            if abs(u[i,j] - u_old) < tolerance:\n",
       "                break\n",
       "\n",
       "# Plot solution\n",
       "X, Y = np.meshgrid(x, y)\n",
       "plt.contourf(X, Y, u.T, 100, cmap='viridis')\n",
       "plt.colorbar()\n",
       "plt.xlabel('x')\n",
       "plt.ylabel('y')\n",
       "plt.title('2D Laplace equation solution')\n",
       "plt.show()\n",
       "```\n",
       "\n",
       "Explanation:\n",
       "\n",
       "This code uses finite difference methods to solve the 2D Laplace equation, which is given by:\n",
       "\n",
       "∇²u(x,y) = 0\n",
       "\n",
       "on the square domain x=(0,1) and y=(0,1) with vanishing boundary conditions. \n",
       "\n",
       "We set up a grid with nx x ny points in the x and y directions, respectively. We initialize a solution vector u with zeros, and then set the boundary conditions to zero. We also set up a source term f, which is a delta function located at the center of the domain. The source term is used to enforce the vanishing boundary conditions.\n",
       "\n",
       "We then use an iterative solver to solve the Laplace equation. The solver updates the values of the solution vector at each point on the grid until the difference between consecutive iterations is below a certain tolerance.\n",
       "\n",
       "Finally, we plot the solution using Matplotlib, which creates a contour plot of the solution on the domain. The contour plot is a visualization of the solution u(x,y) as a height map, which provides a 2D picture of the solution."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%ai chatgpt\n",
    "Please generate the Python code to solve the 2D Laplace equation in cartesian coordinates.\n",
    "Solve the equation on the square domain x=(0,1) and y=(0,1) with vanishing boundary conditions.\n",
    "Plot the solution using Matplotlib.\n",
    "Please also provide an explanation."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
